#include <pinocchio/fwd.hpp>
#include "cartpole_dynamic.h"
#include <ocs2_core/automatic_differentiation/Types.h>
#include <pinocchio/algorithm/aba.hpp>
#include <pinocchio/algorithm/frames.hpp>
namespace cartpole
{
    CartpoleDynamic::CartpoleDynamic(const PinocchioInterface &pinoccho_interface)
    {
        pinocchio_interface_ad_ptr_.reset(new PinocchioInterfaceCppAd(pinoccho_interface.toCppAd()));
        initialize(StateDim, InputDim, "cartpole");
    }
    ad_vector_t CartpoleDynamic::systemFlowMap(ad_scalar_t time, const ad_vector_t &state, const ad_vector_t &input,
                                               const ad_vector_t &parameters) const
    {
        // state    [p \theta1 \theta2 \theta3 v \omega1 \omega2 \theta3]
        // input    [f]
        // q        [p \theta1 \theta2 \theta3]
        // dq       [v \omega1 \omega2 \omega3]
        // input    [f 0 0 0]
        const auto &model = pinocchio_interface_ad_ptr_->getModel();
        auto &data = pinocchio_interface_ad_ptr_->getData();
        // 构建 pinocchio需要的q与dq与 tau
        // using ad_vector_t = Eigen::Matrix<ad_scalar_t, Eigen::Dynamic, 1>;
        ad_vector_t q(model.nq);
        ad_vector_t dq(model.nv);
        ad_vector_t tau = ad_vector_t(model.nv).setZero();
        q = state.head(model.nq);
        dq = state.tail(model.nv);
        tau[0] = input[0];
        ad_vector_t ddq(model.nv);
        ddq = pinocchio::aba(model, data, q, dq, tau);
        ad_vector_t dx(StateDim);
        dx << dq, ddq;
        return dx;
    }

}